import logging
import time
import threading
from config import config, logger

class ModelManager:
    def __init__(self):
        self.models = config.get("models", {})
        self.model_lock = threading.Lock()
        self.loaded_models = {}

    def load_model(self, model_name):
        if model_name not in self.models:
            logger.error(f"Model {model_name} not found")
            return False, "Model not found"

        with self.model_lock:
            if model_name not in self.loaded_models:
                logger.info(f"Loading model {model_name}")
                self.loaded_models[model_name] = {
                    "last_used": time.time(),
                    "config": self.models[model_name]
                }
            else:
                logger.info(f"Model {model_name} is already loaded")
                self.loaded_models[model_name]["last_used"] = time.time()

        return True, "Model loaded successfully"

    def get_model(self, model_name):
        with self.model_lock:
            if model_name not in self.loaded_models:
                logger.warn(f"Model {model_name} not loaded")
                return None
            logger.info(f"Using model {model_name}")
            self.loaded_models[model_name]["last_used"] = time.time()
            return self.loaded_models[model_name]["config"]

    def unload_model(self, model_name):
        with self.model_lock:
            if model_name in self.loaded_models:
                logger.info(f"Unloading model {model_name}")
                del self.loaded_models[model_name]
            else:
                logger.warn(f"Model {model_name} not loaded")

def auto_unload_models():
    model_manager = ModelManager()
    while True:
        current_time = time.time()
        with model_manager.model_lock:
            models_to_unload = [
                model_name for model_name, model_info in model_manager.loaded_models.items()
                if current_time - model_info["last_used"] > 60
            ]
            for model_name in models_to_unload:
                model_manager.unload_model(model_name)
        time.sleep(10)  # Check every 10 seconds
